import random

from driving_gym.environment.scenario.base_scenario import BaseScenario
from driving_gym.simulation.adapter.carla.carla_adapter import CarlaAdapter
from driving_gym.simulation.adapter.carla.carla_config import (
    CarlaActorConfig,
    CarlaSensorConfig,
)
from driving_gym.simulation.adapter.carla.planner import PathTracker


class CarlaScenario(BaseScenario):
    adapter: CarlaAdapter

    def __init__(self, config, adapter):
        super().__init__(config, adapter)

        self.attach_default_sensors = self.config.get("attach_default_sensors", True)
        self.enable_planner = self.config.get("enable_planner", False)
        self.draw_path = self.config.get("draw_path", False)
        self.endless = self.config.get("endless", False)
        self.planners: dict[str, PathTracker] = {}

    def reset(self):
        # Check if any random task is set
        task_config = self.config["task"]
        for actor_id, task in task_config.items():
            if task.get("random", False):
                min_dist = task.get("min_dist", -1)
                max_dist = task.get("max_dist", float("inf"))
                map = self.adapter.map

                while True:
                    wpts = map.generate_waypoints(1)
                    start = random.choice(wpts)
                    end = random.choice(wpts)
                    dist = start.transform.location.distance(end.transform.location)
                    if min_dist < dist < max_dist:
                        task["start"] = start.transform
                        task["end"] = end.transform
                        break

        if self.attach_default_sensors:
            # Add collision_sensor and lane_invasion_sensor for each actor
            for actor_id, actor_config in self.config["actor"].items():
                actor_config["sensors"]["collision"] = (
                    CarlaSensorConfig()
                    .blueprint(type="collision")
                    .spawning(parent=actor_id)
                    .to_dict()
                )
                actor_config["sensors"]["lane_invasion"] = (
                    CarlaSensorConfig()
                    .blueprint(type="lane_invasion")
                    .spawning(parent=actor_id)
                    .to_dict()
                )
        super().reset()

        if self.enable_planner:
            # Setup planner after creating actors
            self._setup_planner()

    def run_step(self):
        if self.enable_planner:
            for actor_id, planner in self.planners.items():
                planner.run_step()
                if self.endless and len(planner.get_planner_path()) < 10:
                    planner.navigate_to_random_location()
                if self.draw_path:
                    planner.draw(life_time=0.5)

    def _setup_planner(self):
        for actor_id, task in self.config["task"].items():
            actor = self.adapter.get_actor(actor_id)
            start = CarlaActorConfig.parse_location(task["start"], self.adapter.map)
            end = CarlaActorConfig.parse_location(task["end"], self.adapter.map)

            start_loc = (start.location.x, start.location.y, start.location.z)
            end_loc = (end.location.x, end.location.y, end.location.z)
            planner = PathTracker(actor, start_loc, end_loc)
            self.planners[actor_id] = planner
